#
# This file is a part of the open-eBackup project.
# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
# If a copy of the MPL was not distributed with this file, You can obtain one at
# http://mozilla.org/MPL/2.0/.
#
# Copyright (c) [2024] Huawei Technologies Co.,Ltd.
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
#

import os
import sys
import urllib.request
import urllib.error
import http.client
import socket
from urllib.response import addinfourl
import ssl
import re

from exchange.commons.logger_wins import Logger
from exchange.ntlm import ntlm_service as ntlm

HEADER = {
    "Content-type": "text/xml; charset=utf-8",
    "Accept": "text/xml",
}

READ_BYTE_LENGTH = 1024 * 1024
LOGGER = Logger().get_logger()
COMPATIBILITY_MODE = False

SECURE_CIPHER_SUITES = ("ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-AES256-GCM-SHA384: \
                         ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305: \
                         DHE-RSA-AES128-GCM-SHA256:DHE-RSA-AES256-GCM-SHA384:DHE-RSA-CHACHA20-POLY1305: \
                         DHE-RSA-AES256-GCM-SHA256:AES128-GCM-SHA256:AES256-GCM-SHA384")

CIPHER_SUITES = ("ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-AES256-GCM-SHA384: \
                  ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305: \
                  DHE-RSA-AES128-GCM-SHA256:DHE-RSA-AES256-GCM-SHA384:DHE-RSA-CHACHA20-POLY1305: \
                  ECDHE-ECDSA-AES128-SHA256:ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES128-SHA: \
                  ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA: \
                  DHE-RSA-AES128-SHA256:DHE-RSA-AES256-SHA256:AES128-GCM-SHA256:AES256-GCM-SHA384:AES128-SHA256: \
                  AES256-SHA256:AES128-SHA:AES256-SHA:DES-CBC3-SHA")

SSL_PROTOCOL_VERSION = ssl.PROTOCOL_TLS


def extract_user_domain_type1_flag(user):
    user_parts = user.split('\\', 1)
    if len(user_parts) == 1:
        user_name = user_parts[0]
        domain_name = ''
        type1_flags = ntlm.NTLM_TYPE1_FLAGS & ~ntlm.NTLMSSP_NEGOTIATE_OEM_DOMAIN_SUPPLIED
    else:
        domain_name = user_parts[0].upper()
        user_name = user_parts[1]
        type1_flags = ntlm.NTLM_TYPE1_FLAGS
    return domain_name, type1_flags, user_name


class AbstractNtlmAuthHandler:

    def __init__(self, password_mgr, file_obj=None):
        if password_mgr is None:
            raise Exception(f"password_mgr cannot be none")
        if file_obj is None:
            self.is_part = False
        else:
            self.is_part = True
            self.file_obj = file_obj
        self.passwd = password_mgr
        self.add_password = self.passwd.add_password

    def http_error_authentication_required(self, auth_header_field, req, fp, headers):
        auth_header_value_list = headers.get_all(auth_header_field)
        if auth_header_value_list:
            if any([hv.lower() == 'ntlm' for hv in auth_header_value_list]):
                fp.close()
                return self.retry_using_http_ntlm_auth(req, auth_header_field, None)
        raise urllib.request.URLError("auth_header_value_list is None")

    def retry_using_http_ntlm_auth(self, req, auth_header_field, realm):
        user, pw = self.passwd.find_user_password(realm, req.get_full_url())
        if pw is not None:
            domain_name, type1_flags, user_name = extract_user_domain_type1_flag(user)
            # ntlm认证过程中必须使用同一个套接字进行完整的握手
            headers = dict(req.headers)
            headers.update(req.unredirected_hdrs)
            auth = 'NTLM %s' % ntlm.create_ntlm_negotiate_message(domain_name, type1_flags)
            if req.headers.get(self.auth_header, None) == auth:
                raise urllib.request.URLError("auth_header is already filled")
            headers[self.auth_header] = auth
            host = req.host
            if not host:
                raise urllib.request.URLError('no host given')
            if req.get_full_url().startswith('https://'):
                http_connection = http.client.HTTPSConnection(host)  # will parse host:port
            else:
                http_connection = http.client.HTTPConnection(host)  # will parse host:port
            # we must keep the connection because NTLM authenticates the connection, not single requests
            headers["Connection"] = "Keep-Alive"
            headers = dict((name.title(), val) for name, val in list(headers.items()))
            http_connection.request(req.get_method(), req.selector, req.data, headers)
            response = http_connection.getresponse()
            response.begin()
            response.read(int(response.getheader('content-length')))
            try:
                if response.getheader('set-cookie'):
                    # 存cookie
                    headers['Cookie'] = response.getheader('set-cookie')
            except TypeError:
                pass
            # 移除对套接字的引用，这样它就不能被响应对象关闭（我们希望保持套接字打开）
            response.fp = None
            auth_header_value = response.getheader(auth_header_field, None)
            # 某些Exchange服务器发送两个WWW-Authenticate报头，一个带有NTLM challenge，另一个带有'Negotiate'关键字 - 确保获取正确的请求头
            is_match = re.match('(NTLM [A-Za-z0-9+\-/=]+)', auth_header_value)
            if is_match:
                auth_header_value, = is_match.groups()
            server_challenge, negotiate_flags = ntlm.parse_ntlm_challenge_message(auth_header_value[5:])
            auth = 'NTLM %s' % ntlm.create_ntlm_authenticate_message(
                server_challenge, user_name, domain_name, pw, negotiate_flags
            )
            headers[self.auth_header] = auth
            headers["Connection"] = "Close"
            headers = dict((name.title(), val) for name, val in list(headers.items()))
            if self.is_part:
                headers['Content-Length'] = self.file_obj.total_byte_len
            try:
                response = self.send_data_and_get_response(headers, http_connection, req, response)

                def notimplemented():
                    raise NotImplementedError

                response.readline = notimplemented
                return addinfourl(response, response.msg, req.get_full_url(), response.code)
            except socket.error as err:
                raise urllib.request.URLError(err)
        else:
            raise urllib.request.URLError("pw is None")

    def send_data_and_get_response(self, headers, http_connection, req, response):
        if self.is_part:
            _real_send_bytes = 0
            http_connection.request(req.get_method(), req.selector, b'', headers)
            while True:
                byte_data = self.file_obj.read(READ_BYTE_LENGTH)
                if len(byte_data) > 0:
                    http_connection.send(byte_data)
                else:
                    break
        else:
            http_connection.request(req.get_method(), req.selector, req.data, headers)
        # 重定向请求无法处理
        response = http_connection.getresponse()
        return response


class HTTPNtlmAuthHandler(AbstractNtlmAuthHandler, urllib.request.BaseHandler):
    auth_header = 'Authorization'

    def http_error_401(self, req, fp, code, msg, headers):
        return self.http_error_authentication_required('www-authenticate', req, fp, headers)


def set_ssl_protocol_version(windows_version: str):
    # 6.2代表 Windows server 2012, 只支持到TLS 1.1
    if windows_version.startswith("6.2") or windows_version.startswith("6.3"):
        global SSL_PROTOCOL_VERSION
        SSL_PROTOCOL_VERSION = ssl.PROTOCOL_TLSv1_1
        LOGGER.debug(f"get version: {windows_version} and user tls: {SSL_PROTOCOL_VERSION}")


def get_post_response_stream(url: str, post_body: str, username: str, password: str, **kwargs):
    global COMPATIBILITY_MODE
    if kwargs.get("compatibility_mode") == "true":
        COMPATIBILITY_MODE = True
    file_obj = kwargs.get("file_obj", None)
    windows_version = kwargs.get("windows_version", "")
    set_ssl_protocol_version(windows_version)
    pass_mgr = urllib.request.HTTPPasswordMgrWithDefaultRealm()
    pass_mgr.add_password(None, url, username, password)
    auth_ntlm_handler = HTTPNtlmAuthHandler(pass_mgr, file_obj)
    proxy_handler = urllib.request.ProxyHandler({})
    # 全局禁用urllib HTTPS证书认证，由于插件python进程不是常驻的，所以此处修改不会影响到其他备份任务
    ssl._create_default_https_context = create_unverified_context
    opener = urllib.request.build_opener(proxy_handler, auth_ntlm_handler)
    urllib.request.install_opener(opener)
    req = urllib.request.Request(url=url, headers=HEADER)
    response = urllib.request.urlopen(req, data=post_body.encode())
    return response


def create_unverified_context(protocol=None, *, cert_reqs=ssl.CERT_NONE,
                              check_hostname=False, purpose=ssl.Purpose.SERVER_AUTH,
                              certfile=None, keyfile=None,
                              cafile=None, capath=None, cadata=None):
    # 尝试使用所有的协议版本
    context = ssl.SSLContext(SSL_PROTOCOL_VERSION)
    if COMPATIBILITY_MODE:
        context.set_ciphers(CIPHER_SUITES)
    else:
        context.set_ciphers(SECURE_CIPHER_SUITES)
    context.check_hostname = check_hostname
    if cert_reqs is not None:
        context.verify_mode = cert_reqs
    if check_hostname:
        context.check_hostname = True

    if keyfile and not certfile:
        raise ValueError("certfile must be specified")
    if certfile or keyfile:
        context.load_cert_chain(certfile, keyfile)

    # load CA root certs
    if cafile or capath or cadata:
        context.load_verify_locations(cafile, capath, cadata)
    elif context.verify_mode != ssl.CERT_NONE:
        # no explicit cafile, capath or cadata but the verify mode is
        # CERT_OPTIONAL or CERT_REQUIRED. Let's try to load default system
        # root CA certificates for the given purpose. This may fail silently.
        context.load_default_certs(purpose)
    # OpenSSL 1.1.1 keylog file
    if hasattr(context, 'keylog_filename'):
        keylogfile = os.environ.get('SSLKEYLOGFILE')
        if keylogfile and not sys.flags.ignore_environment:
            context.keylog_filename = keylogfile
    return context
